{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Agent Conversation with A2A\n",
    "\n",
    "This notebook demonstrates how to use the Python A2A package to create conversations between agents. We'll explore:\n",
    "\n",
    "1. Creating A2A messages\n",
    "2. Building conversations\n",
    "3. Exchanging messages between agents\n",
    "4. Working with different message types\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's make sure we have the Python A2A package installed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install python-a2a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's import the necessary components from the package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from python_a2a import (\n",
    "    Message, TextContent, FunctionCallContent, FunctionResponseContent, FunctionParameter,\n",
    "    MessageRole, Conversation,\n",
    "    A2AClient, A2AServer,\n",
    "    pretty_print_message, pretty_print_conversation\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Messages\n",
    "\n",
    "Let's start by creating a simple text message:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a text message from the user\n",
    "text_message = Message(\n",
    "    content=TextContent(text=\"What's the weather like in New York?\"),\n",
    "    role=MessageRole.USER\n",
    ")\n",
    "\n",
    "# Display the message\n",
    "pretty_print_message(text_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also create function call messages, which are used to request specific actions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a function call message\n",
    "function_call = Message(\n",
    "    content=FunctionCallContent(\n",
    "        name=\"get_weather\",\n",
    "        parameters=[\n",
    "            FunctionParameter(name=\"location\", value=\"New York\"),\n",
    "            FunctionParameter(name=\"unit\", value=\"celsius\")\n",
    "        ]\n",
    "    ),\n",
    "    role=MessageRole.AGENT\n",
    ")\n",
    "\n",
    "# Display the function call\n",
    "pretty_print_message(function_call)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And function response messages, which return the results of a function call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a function response message\n",
    "function_response = Message(\n",
    "    content=FunctionResponseContent(\n",
    "        name=\"get_weather\",\n",
    "        response={\n",
    "            \"temperature\": 22,\n",
    "            \"conditions\": \"Partly Cloudy\",\n",
    "            \"humidity\": 65,\n",
    "            \"wind_speed\": 8\n",
    "        }\n",
    "    ),\n",
    "    role=MessageRole.AGENT,\n",
    "    parent_message_id=function_call.message_id\n",
    ")\n",
    "\n",
    "# Display the function response\n",
    "pretty_print_message(function_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Conversations\n",
    "\n",
    "Now let's build a conversation that includes multiple messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new conversation\n",
    "conversation = Conversation()\n",
    "\n",
    "# Add messages to the conversation\n",
    "conversation.add_message(text_message)\n",
    "conversation.add_message(function_call)\n",
    "conversation.add_message(function_response)\n",
    "\n",
    "# Add an agent's text response\n",
    "weather_response = Message(\n",
    "    content=TextContent(\n",
    "        text=\"The weather in New York is currently 22°C (71.6°F) and partly cloudy. \"\n",
    "             \"The humidity is 65% with wind speeds of 8 km/h.\"\n",
    "    ),\n",
    "    role=MessageRole.AGENT,\n",
    "    parent_message_id=function_response.message_id,\n",
    "    conversation_id=conversation.conversation_id\n",
    ")\n",
    "conversation.add_message(weather_response)\n",
    "\n",
    "# Display the conversation\n",
    "pretty_print_conversation(conversation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use helper methods to create messages directly in the conversation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new conversation\n",
    "conversation2 = Conversation()\n",
    "\n",
    "# Add a user message\n",
    "msg1 = conversation2.create_text_message(\n",
    "    text=\"What's the weather like in New York?\",\n",
    "    role=MessageRole.USER\n",
    ")\n",
    "\n",
    "# Add a function call\n",
    "msg2 = conversation2.create_function_call(\n",
    "    name=\"get_weather\",\n",
    "    parameters=[\n",
    "        {\"name\": \"location\", \"value\": \"New York\"},\n",
    "        {\"name\": \"unit\", \"value\": \"celsius\"}\n",
    "    ],\n",
    "    role=MessageRole.AGENT,\n",
    "    parent_message_id=msg1.message_id\n",
    ")\n",
    "\n",
    "# Add a function response\n",
    "msg3 = conversation2.create_function_response(\n",
    "    name=\"get_weather\",\n",
    "    response={\n",
    "        \"temperature\": 22,\n",
    "        \"conditions\": \"Partly Cloudy\",\n",
    "        \"humidity\": 65,\n",
    "        \"wind_speed\": 8\n",
    "    },\n",
    "    role=MessageRole.AGENT,\n",
    "    parent_message_id=msg2.message_id\n",
    ")\n",
    "\n",
    "# Add an agent's text response\n",
    "msg4 = conversation2.create_text_message(\n",
    "    text=\"The weather in New York is currently 22°C (71.6°F) and partly cloudy. \"\n",
    "         \"The humidity is 65% with wind speeds of 8 km/h.\",\n",
    "    role=MessageRole.AGENT,\n",
    "    parent_message_id=msg3.message_id\n",
    ")\n",
    "\n",
    "# Display the conversation\n",
    "pretty_print_conversation(conversation2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serializing and Deserializing Messages\n",
    "\n",
    "A2A messages and conversations can be serialized to JSON for transport and storage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Serialize a message to JSON\n",
    "json_message = text_message.to_json()\n",
    "print(json_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deserialize a message from JSON\n",
    "deserialized_message = Message.from_json(json_message)\n",
    "pretty_print_message(deserialized_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Serialize a conversation to JSON\n",
    "json_conversation = conversation.to_json()\n",
    "print(json_conversation[:500] + \"...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deserialize a conversation from JSON\n",
    "deserialized_conversation = Conversation.from_json(json_conversation)\n",
    "pretty_print_conversation(deserialized_conversation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a Simple Agent Server\n",
    "\n",
    "Let's define a simple echo agent that responds to messages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EchoAgent(A2AServer):\n",
    "    \"\"\"A simple agent that echoes back messages.\"\"\"\n",
    "    \n",
    "    def handle_message(self, message):\n",
    "        \"\"\"Process incoming A2A messages.\"\"\"\n",
    "        if message.content.type == \"text\":\n",
    "            return Message(\n",
    "                content=TextContent(text=f\"Echo: {message.content.text}\"),\n",
    "                role=MessageRole.AGENT,\n",
    "                parent_message_id=message.message_id,\n",
    "                conversation_id=message.conversation_id\n",
    "            )\n",
    "        else:\n",
    "            return Message(\n",
    "                content=TextContent(text=f\"Received a {message.content.type} message type\"),\n",
    "                role=MessageRole.AGENT,\n",
    "                parent_message_id=message.message_id,\n",
    "                conversation_id=message.conversation_id\n",
    "            )\n",
    "\n",
    "# Create an instance of the echo agent\n",
    "echo_agent = EchoAgent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's send a message to our echo agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a test message\n",
    "test_message = Message(\n",
    "    content=TextContent(text=\"Hello, agent!\"),\n",
    "    role=MessageRole.USER\n",
    ")\n",
    "\n",
    "# Send the message to the echo agent\n",
    "response = echo_agent.handle_message(test_message)\n",
    "\n",
    "# Display the response\n",
    "pretty_print_message(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also send a conversation to the agent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new conversation\n",
    "test_conversation = Conversation()\n",
    "test_conversation.create_text_message(\n",
    "    text=\"Hello, agent!\",\n",
    "    role=MessageRole.USER\n",
    ")\n",
    "\n",
    "# Send the conversation to the echo agent\n",
    "updated_conversation = echo_agent.handle_conversation(test_conversation)\n",
    "\n",
    "# Display the updated conversation\n",
    "pretty_print_conversation(updated_conversation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running a Server and Connecting with a Client\n",
    "\n",
    "In a real application, you would run the agent as a server using `run_server()`, and then connect to it using `A2AClient`. Here's how you would do that if the server was running:\n",
    "\n",
    "```python\n",
    "# Run the server (in a separate process)\n",
    "run_server(echo_agent, host=\"0.0.0.0\", port=5000)\n",
    "\n",
    "# Connect with a client\n",
    "client = A2AClient(\"http://localhost:5000/a2a\")\n",
    "\n",
    "# Send a message\n",
    "response = client.send_message(test_message)\n",
    "```\n",
    "\n",
    "For this notebook, we'll just simulate the client-server interaction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate a client-server interaction\n",
    "def simulate_client_server(agent, message):\n",
    "    \"\"\"Simulate a client sending a message to a server.\"\"\"\n",
    "    print(\"Client sending message:\")\n",
    "    pretty_print_message(message)\n",
    "    \n",
    "    print(\"\\nServer processing message...\\n\")\n",
    "    response = agent.handle_message(message)\n",
    "    \n",
    "    print(\"Client received response:\")\n",
    "    pretty_print_message(response)\n",
    "    \n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate a client interaction with our echo agent\n",
    "client_message = Message(\n",
    "    content=TextContent(text=\"What is the meaning of life?\"),\n",
    "    role=MessageRole.USER\n",
    ")\n",
    "\n",
    "server_response = simulate_client_server(echo_agent, client_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we've learned how to:\n",
    "\n",
    "1. Create different types of A2A messages (text, function calls, function responses)\n",
    "2. Build and manage conversations\n",
    "3. Serialize and deserialize messages and conversations\n",
    "4. Create a simple agent that can handle messages\n",
    "5. Simulate client-server interactions\n",
    "\n",
    "The A2A protocol provides a standardized way for agents to communicate, making it easy to build modular AI systems where specialized agents can collaborate to solve complex problems."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}